import os
import subprocess
from multiprocessing import Pool



dirs = [
    '/mnt/nas/zhangshu/train_data/爬取_clean',        # 2056
    '/mnt/nas/zhangshu/train_data/igpair/clothes',    # 86873
    '/mnt/nas/zhangshu/train_data/h-and-m-fashion-caption-12k/images',  # 12437
    '/mnt/nas/zhangshu/train_data/dresscode/images',  # 44263
    '/mnt/nas/zhangshu/train_data/viton/images',       # 11647
    '/mnt/nas/shengjie/datasets/cloth_all_images/', # 15xxx
]

def get_disk_usage(path):
    """Calculate disk usage for a directory using du -sh"""
    try:
        result = subprocess.run(['du', '-sh', path], 
                              stdout=subprocess.PIPE, 
                              stderr=subprocess.PIPE,
                              text=True)
        if result.returncode == 0:
            size = result.stdout.split()[0]
            return (path, size)
        else:
            return (path, f"Error: {result.stderr.strip()}")
    except Exception as e:
        return (path, f"Exception: {str(e)}")

def main():
    print("Calculating disk usage for directories:")
    print("-" * 50)
    
    # Using multiprocessing to speed up the process
    with Pool(processes=len(dirs)) as pool:
        results = pool.map(get_disk_usage, dirs)
    
    # Print results in a formatted table
    max_path_len = max(len(path) for path, _ in results)
    print(f"{'Directory Path':<{max_path_len}} | Size")
    print("-" * (max_path_len + 3 + 10))
    
    for path, size in results:
        print(f"{path:<{max_path_len}} | {size}")

if __name__ == '__main__':
    main()